import numpy as np

from ultralytics import YOLO

classes = {0: "half-ripe", 1: "raw", 2: "ripe"}

# 1、加载模型
# 加载YOLOv8模型
model = YOLO("./runs/detect/train/weights/best.pt")
model(np.random.rand(100, 100, 3))

def predict_yolo(image):
    # 在帧上运行YOLOv8推理
    results = model(image)
    classes = results[0].boxes.cls
    # 如果有类别名称，可以通过类别索引获取
    class_names = model.names[int(classes)]
    return class_names
"""
此模块定义了用于货物分类的YOLOv8推理函数，供主分拣系统调用。

用途：提供predict_yolo函数，将输入图像分类为“成熟”（ripe）、“半成熟”（half-ripe）或“未成熟”（raw）。
核心组件：

类别映射：定义了YOLO类别索引到标签的映射：
pythonclasses = {0: "half-ripe", 1: "raw", 2: "ripe"}

模型加载：加载自定义训练的YOLOv8模型（./runs/detect/train/weights/best.pt）。通过对随机图像进行一次推理（model(np.random.rand(100, 100, 3))）初始化模型。
推理函数（predict_yolo）：

接受输入图像（OpenCV的numpy数组）。
运行YOLOv8推理（results = model(image)）。
从检测结果中提取类别索引（results[0].boxes.cls）。
使用model.names[int(classes)]将类别索引映射为类别名称。
返回类别名称（“ripe”、“half-ripe”或“raw”）。




用途：被main.py和test.py导入，用于对通过MQTT接收的货物图像进行分类。
"""